# -----------------------------------------------------------------------#
#   predict.py将单张图片预测、摄像头检测、FPS测试和目录遍历检测等功能
#   整合到了一个py文件中，通过指定mode进行模式的修改。
# -----------------------------------------------------------------------#
import time
import math
import cv2
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
from yolo import YOLO
import os
from deeplab_v3_test import deeplab_detect
from matplotlib.gridspec import GridSpec

if __name__ == "__main__":
    yolo = YOLO()
    # ----------------------------------------------------------------------------------------------------------#
    #   mode用于指定测试的模式：
    #   'predict'表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释
    #   'video'表示视频检测，可调用摄像头或者视频进行检测，详情查看下方注释。
    #   'fps'表示测试fps，使用的图片是img里面的street.jpg，详情查看下方注释。
    #   'dir_predict'表示遍历文件夹进行检测并保存。默认遍历img文件夹，保存img_out文件夹，详情查看下方注释。
    # ----------------------------------------------------------------------------------------------------------#
    mode = "predict"
    # ----------------------------------------------------------------------------------------------------------#
    #   video_path用于指定视频的路径，当video_path=0时表示检测摄像头
    #   想要检测视频，则设置如video_path = "xxx.mp4"即可，代表读取出根目录下的xxx.mp4文件。
    #   video_save_path表示视频保存的路径，当video_save_path=""时表示不保存
    #   想要保存视频，则设置如video_save_path = "yyy.mp4"即可，代表保存为根目录下的yyy.mp4文件。
    #   video_fps用于保存的视频的fps
    #   video_path、video_save_path和video_fps仅在mode='video'时有效
    #   保存视频时需要ctrl+c退出或者运行到最后一帧才会完成完整的保存步骤。
    # ----------------------------------------------------------------------------------------------------------#
    video_path = r'C:\Users\jiangzhengquan\Desktop\video\磨煤机\4.mp4'
    mask_dir = r'C:\Users\jiangzhengquan\Desktop\Switch_Dataset\SegmentationClass'
    save_dir = r'C:\Users\jiangzhengquan\Desktop\Switch_Dataset\SmallImages'
    result_dir = r'C:\Users\jiangzhengquan\Desktop\Switch_Dataset\result'
    video_save_path = ""
    video_fps = 25.0
    # -------------------------------------------------------------------------#
    #   test_interval用于指定测量fps的时候，图片检测的次数
    #   理论上test_interval越大，fps越准确。
    # -------------------------------------------------------------------------#
    test_interval = 100
    # -------------------------------------------------------------------------#
    #   dir_origin_path指定了用于检测的图片的文件夹路径
    #   dir_save_path指定了检测完图片的保存路径
    #   dir_origin_path和dir_save_path仅在mode='dir_predict'时有效
    # -------------------------------------------------------------------------#
    dir_origin_path = r"C:\Users\jiangzhengquan\Desktop\Switch_Dataset\JPEGImages"
    dir_save_path = "img_out/"

    if mode == "predict":
        '''
        1、如果想要进行检测完的图片的保存，利用r_image.save("img.jpg")即可保存，直接在predict.py里进行修改即可。 
        2、如果想要获得预测框的坐标，可以进入yolo.detect_image函数，在绘图部分读取top，left，bottom，right这四个值。
        3、如果想要利用预测框截取下目标，可以进入yolo.detect_image函数，在绘图部分利用获取到的top，left，bottom，right这四个值
        在原图上利用矩阵的方式进行截取。
        4、如果想要在预测图上写额外的字，比如检测到的特定目标的数量，可以进入yolo.detect_image函数，在绘图部分对predicted_class进行判断，
        比如判断if predicted_class == 'car': 即可判断当前目标是否为车，然后记录数量即可。利用draw.text即可写字。
        '''
        img_name = '3_1.jpg'
        image_path = os.path.join(dir_origin_path, img_name)
        src_img = cv2.imread(image_path)
        image = Image.open(image_path)

        r_image, box,labels = yolo.detect_image(image)

        img_name = os.path.splitext(img_name)[0]
        mask_path = os.path.join(mask_dir, img_name + '.png')
        mask_img = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        h, w = src_img.shape[:2]
        box = box.astype(np.int32)
        print(img_name)
        for i, (y1, x1, y2, x2) in enumerate(box):
            x1 = max(x1 - 30, 1)
            x2 = min(x2 + 100, w)
            y1 = max(y1 - 30, 1)
            y2 = min(y2 + 100, h)
            src_save_img = src_img[y1:y2, x1:x2, :]
            mask_save_img = mask_img[y1:y2, x1:x2]

            plt.imsave(os.path.join(save_dir, img_name + "_" + str(i) + '.jpg'), src_save_img)
            cv2.imwrite(os.path.join(save_dir, img_name + "_" + str(i) + '.png'), mask_save_img)
            fp.write(img_name + "_" + str(i) + '\n')

    elif mode == "video":
        capture = cv2.VideoCapture(video_path)
        if video_save_path != "":
            fourcc = cv2.VideoWriter_fourcc(*'I420')
            size = (513, 513)
            video_save_path = os.path.join(video_save_path, "small_" + video_path.rsplit('\\')[-1])
            out = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)

        fps = 0.0
        fig = plt.figure(figsize=(20, 20))
        gs = GridSpec(20, 40)
        count = 0
        while (True):
            t1 = time.time()
            # 读取某一帧
            ref, image = capture.read()
            if image.shape[0] != 1080:
                image = cv2.resize(image, (1920, 1080))
            if not ref:
                break
            # 格式转变，BGRtoRGB
            # frame = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            # 转变成Image
            frame = Image.fromarray(np.uint8(image))
            # 进行检测
            frame = yolo.detect_image(frame)
            # frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            # cv2.imshow('', frame)
            # cv2.waitKey(1)
            # cv2.imshow("video", frame)
            if isinstance(frame, tuple):
                frame, boxes, labels = frame
                fps = (fps + (1. / (time.time() - t1))) / 2
                # RGBtoBGR满足opencv显示格式
                frame = np.asarray(frame)
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

                frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

                # cv2.imshow("video", frame)
                # c = cv2.waitKey(1) & 0xff
                # if video_save_path != "":
                h, w = image.shape[:2]
                origins = []
                for label, box in list(zip(labels, boxes)):
                    y1, x1, y2, x2 = list(map(int, box))
                    x1 = max(x1 - 30, 0)
                    x2 = min(x2 + 100, w)
                    y1 = max(y1 - 30, 0)
                    y2 = min(y2 + 100, h)
                    save_img = image[y1:y2, x1:x2, :]
                    origin = deeplab_detect(save_img, label)
                    origins.append(origin)
                fps = (fps + (1. / (time.time() - t1))) / 2
                print("fps= %.2f" % (fps))
                plt.ion()
                # ax1 = plt.subplot(1, 2, 1)
                # ax2 = plt.subplot(1, 2, 2)
                length = 20 // len(origins)

                ax1 = fig.add_subplot(gs[:, 0:30])
                ax1.imshow(image[:, :, ::-1])
                for i, o in enumerate(origins):
                    ax2 = fig.add_subplot(gs[i * length:(i + 1) * length, 30:])
                    ax2.imshow(o[:, :, ::-1])
                plt.pause(0.01)
                plt.show()
                plt.savefig(
                    os.path.join(result_dir, video_path.rsplit('\\')[-1].split('.')[0] + '_' + str(count) + '.jpg'))
                plt.clf()
                count += 1

            # if c == 27:
            #     capture.release()
            #     break
        capture.release()
        # out.release()
        cv2.destroyAllWindows()

    elif mode == "fps":
        img = Image.open('img/street.jpg')
        tact_time = yolo.get_FPS(img, test_interval)
        print(str(tact_time) + ' seconds, ' + str(1 / tact_time) + 'FPS, @batch_size 1')

    elif mode == "dir_predict":
        import os
        from tqdm import tqdm

        img_names = os.listdir(dir_origin_path)
        with open(os.path.join(save_dir, 'train.txt'), 'w') as fp:
            for img_name in tqdm(img_names):
                if img_name.lower().endswith(
                        ('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
                    image_path = os.path.join(dir_origin_path, img_name)
                    src_img = cv2.imread(image_path)
                    image = Image.open(image_path)
                    r_image, box = yolo.detect_image(image)

                    img_name = os.path.splitext(img_name)[0]
                    mask_path = os.path.join(mask_dir, img_name + '.png')
                    mask_img = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
                    h, w = src_img.shape[:2]
                    box = box.astype(np.int32)
                    print(img_name)
                    for i, (y1, x1, y2, x2) in enumerate(box):
                        x1 = max(x1 - 30, 1)
                        x2 = min(x2 + 100, w)
                        y1 = max(y1 - 30, 1)
                        y2 = min(y2 + 100, h)
                        src_save_img = src_img[y1:y2, x1:x2, :]
                        mask_save_img = mask_img[y1:y2, x1:x2]

                        plt.imsave(os.path.join(save_dir, img_name + "_" + str(i) + '.jpg'), src_save_img)
                        cv2.imwrite(os.path.join(save_dir, img_name + "_" + str(i) + '.png'), mask_save_img)
                        fp.write(img_name + "_" + str(i) + '\n')

    else:
        raise AssertionError("Please specify the correct mode: 'predict', 'video', 'fps' or 'dir_predict'.")
